"""
document_cleaner 模块
"""

"""
document_cleaner 模块
负责清洗文档文本（去除噪声、规范格式），为分块和嵌入做准备
"""
from abc import ABC, abstractmethod
from typing import List, Dict, Any
import re
import string


class BaseDocumentCleaner(ABC):
    """文档清洗器基类"""

    @abstractmethod
    def clean_text(self, text: str) -> str:
        """
        清洗单段文本

        Args:
            text: 原始文本

        Returns:
            清洗后的文本
        """
        pass

    @abstractmethod
    def clean_documents(self, documents: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """
        清洗文档列表（保留元数据）

        Args:
            documents: 文档列表，每个文档包含 "page_content" 和 "metadata"

        Returns:
            清洗后的文档列表
        """
        pass


class DefaultDocumentCleaner(BaseDocumentCleaner):
    """默认文档清洗器，适用于大多数文本类型"""

    def __init__(self):
        # 定义需要清洗的模式（可根据需求扩展）
        self.patterns = {
            # 去除多余空行（保留单空行）
            "extra_newlines": re.compile(r"\n{3,}"),
            # 去除连续空格（保留单个空格）
            "extra_spaces": re.compile(r" +"),
            # 去除制表符和垂直制表符
            "tabs": re.compile(r"[\t\v]"),
            # 去除无意义的特殊字符（保留中英文标点和字母数字）
            "invalid_chars": re.compile(r"[^\u4e00-\u9fa5a-zA-Z0-9，。,./;:'\"()（）《》【】<>？?！!；;：:]"),
            # 去除首尾空格和换行
            "leading_trailing": re.compile(r"^\s+|\s+$")
        }

    def clean_text(self, text: str) -> str:
        """核心清洗逻辑：按顺序应用清洗规则"""
        if not text:
            return ""

        # 1. 去除制表符
        cleaned = self.patterns["tabs"].sub(" ", text)
        # 2. 去除多余空行（替换为两个换行）
        cleaned = self.patterns["extra_newlines"].sub("\n\n", cleaned)
        # 3. 去除多余空格
        cleaned = self.patterns["extra_spaces"].sub(" ", cleaned)
        # 4. 去除无效特殊字符
        cleaned = self.patterns["invalid_chars"].sub("", cleaned)
        # 5. 去除首尾空白
        cleaned = self.patterns["leading_trailing"].sub("", cleaned)

        return cleaned

    def clean_documents(self, documents: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """批量清洗文档列表，保留元数据"""
        cleaned_docs = []
        for doc in documents:
            cleaned_content = self.clean_text(doc.get("page_content", ""))
            # 仅保留非空文本（过滤完全空白的文档）
            if cleaned_content:
                cleaned_docs.append({
                    "page_content": cleaned_content,
                    "metadata": doc.get("metadata", {})  # 保留原始元数据
                })
        return cleaned_docs


class DocumentCleanerFactory:
    """文档清洗器工厂类，支持扩展不同清洗策略"""

    @staticmethod
    def get_cleaner(cleaner_type: str = "default") -> BaseDocumentCleaner:
        """根据类型获取清洗器实例"""
        if cleaner_type == "default":
            return DefaultDocumentCleaner()
        # 未来可扩展其他清洗器（如针对PDF的特殊清洗器）
        raise ValueError(f"不支持的清洗器类型: {cleaner_type}")